{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1e4049b7-41f6-4b42-bf47-d3f544a2bed9",
   "metadata": {},
   "source": [
    "# Fine-tuning Llama 3.2 with Vision Capabilities - Model Training and Inference\n",
    "\n",
    "## Introduction\n",
    "\n",
    "In this notebook, we'll use the data prepared in the previous notebook to fine-tune a Llama 3.2 multi-modal model using Amazon Bedrock. After fine-tuning, we'll test the model's performance using the test dataset.\n",
    "\n",
    "## Setup\n",
    "\n",
    "First, let's install and import the necessary libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28e2e339-34a5-4e2b-a566-198d05bf74b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required libraries\n",
    "%pip install --upgrade pip\n",
    "%pip install boto3 pillow tqdm matplotlib --upgrade --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f51ebfc-3320-419c-9206-44e66bfe8eaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Restart kernel to ensure updated packages take effect\n",
    "from IPython.core.display import HTML\n",
    "HTML(\"<script>Jupyter.notebook.kernel.restart()</script>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a36ab4a7-1068-4f2c-8c60-b440e4569f3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import os\n",
    "import json\n",
    "import time\n",
    "import base64\n",
    "import io\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "from tqdm.notebook import tqdm\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8c0f1f0-ddb1-45ce-bab1-665f2b8bacfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set AWS region\n",
    "region = \"us-west-2\"  # Llama 3.2 fine-tuning is only available in us-west-2\n",
    "\n",
    "# Create AWS clients\n",
    "session = boto3.session.Session(region_name=region)\n",
    "s3_client = session.client('s3')\n",
    "bedrock = session.client(service_name=\"bedrock\", region_name=region)\n",
    "bedrock_runtime = session.client(service_name=\"bedrock-runtime\", region_name=region)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00890e32-562c-4420-a504-0016dfb73441",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve stored variables from previous notebook\n",
    "%store -r bucket_name\n",
    "%store -r train_data_uri\n",
    "%store -r validation_data_uri\n",
    "%store -r test_data_uri\n",
    "%store -r role_arn\n",
    "%store -r role_name\n",
    "%store -r policy_arn\n",
    "\n",
    "print(f\"Bucket name: {bucket_name}\")\n",
    "print(f\"Training data URI: {train_data_uri}\")\n",
    "print(f\"Validation data URI: {validation_data_uri}\")\n",
    "print(f\"Role ARN: {role_arn}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13d4c8c3-4170-4fd1-af2c-bfa2271433a3",
   "metadata": {},
   "source": [
    "## Create Fine-tuning Job\n",
    "\n",
    "Now, we'll create a fine-tuning job for the Llama 3.2 multi-modal model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9e00e12-bf0a-43af-aed6-77420b3b0114",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a timestamp for unique naming\n",
    "timestamp = time.strftime(\"%Y-%m-%d-%H-%M-%S\")\n",
    "\n",
    "# Define job parameters\n",
    "job_name = f\"llama32-multimodal-ft-{timestamp}\"\n",
    "custom_model_name = f\"llama32-multimodel-{timestamp}\"\n",
    "base_model_id = \"meta.llama3-2-90b-instruct-v1:0:128k\"  # Llama 3.2 vision model ID\n",
    "\n",
    "# Define hyperparameters\n",
    "hyperparameters = {\n",
    "    \"epochCount\": \"2\",       # Number of training epochs\n",
    "    \"batchSize\": \"1\",        # Batch size for training\n",
    "    \"learningRate\": \"0.00001\"  # Learning rate\n",
    "}\n",
    "\n",
    "# Define output location\n",
    "output_s3_uri = f\"s3://{bucket_name}/output/\"\n",
    "\n",
    "# Create validation data config\n",
    "validation_data_config = {\n",
    "    \"validators\": [{\n",
    "        \"s3Uri\": validation_data_uri\n",
    "    }]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "014a95dc-ad46-4bbf-aef5-335e6cf400ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create fine-tuning job\n",
    "try:\n",
    "    response = bedrock.create_model_customization_job(\n",
    "        customizationType=\"FINE_TUNING\",\n",
    "        jobName=job_name,\n",
    "        customModelName=custom_model_name,\n",
    "        roleArn=role_arn,\n",
    "        baseModelIdentifier=base_model_id,\n",
    "        hyperParameters=hyperparameters,\n",
    "        trainingDataConfig={\"s3Uri\": train_data_uri},\n",
    "        validationDataConfig=validation_data_config,\n",
    "        outputDataConfig={\"s3Uri\": output_s3_uri}\n",
    "    )\n",
    "    \n",
    "    # Get job identifier\n",
    "    job_arn = response[\"jobArn\"]\n",
    "    print(f\"Fine-tuning job created: {job_arn}\")\n",
    "    \n",
    "except Exception as e:\n",
    "    print(f\"Error creating fine-tuning job: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fce36b0-f24c-42be-be51-f0757ea719ec",
   "metadata": {},
   "source": [
    "## Monitor Job Status\n",
    "\n",
    "Let's monitor the status of our fine-tuning job:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbb11476-77b5-4e36-8b9d-9f9b1312372d",
   "metadata": {},
   "source": [
    "<div style=\"\n",
    "    background-color: #fcf8e3; \n",
    "    color: #8a6d3b;\n",
    "    padding: 15px;\n",
    "    margin-bottom: 20px;\n",
    "    border: 1px solid #faebcc;\n",
    "    border-radius: 4px;\n",
    "    font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;\">\n",
    "    <span style=\"font-weight:bold;\">⚠️ Warning:</span> \n",
    "    <p>Fine-tuning jobs for Llama 3.2 multi-modal models may take <b>several hours to complete</b>. \n",
    "    The exact duration depends on your dataset size, model parameters, and current training resource availability.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12f3ba73-ce64-447e-862d-99b81b249f34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to check job status\n",
    "def check_job_status(job_arn):\n",
    "    response = bedrock.get_model_customization_job(jobIdentifier=job_arn)\n",
    "    return response[\"status\"]\n",
    "\n",
    "# Get current job status\n",
    "current_status = check_job_status(job_arn)\n",
    "print(f\"Current job status: {current_status}\")\n",
    "\n",
    "# If job completed successfully, get the model details\n",
    "if current_status == \"Completed\":\n",
    "    model_details = bedrock.get_model_customization_job(jobIdentifier=job_arn)\n",
    "    custom_model_arn = model_details[\"outputModelArn\"]\n",
    "    print(f\"Fine-tuned model ARN: {custom_model_arn}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1b85a79-87df-4ecc-bde0-922bd1d6c878",
   "metadata": {},
   "source": [
    "<div style=\"\n",
    "    background-color: #fcf8e3; \n",
    "    color: #8a6d3b;\n",
    "    padding: 15px;\n",
    "    margin-bottom: 20px;\n",
    "    border: 1px solid #faebcc;\n",
    "    border-radius: 4px;\n",
    "    font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;\">\n",
    "    <span style=\"font-weight:bold;\">⚠️ Warning:</span> \n",
    "    <p>Please ensure the status is <b>\"Completed\"</b> before proceeding with the following cells. \n",
    "    You can re-run the status check cell above to update this status.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f4b83ca-3f73-4224-9b9f-4affc8eda97f",
   "metadata": {},
   "source": [
    "## Visualize Training Metrics\n",
    "\n",
    "Let's download and visualize the training metrics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36582331-ae5a-4656-ba62-75c82bd52a0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download training metrics from S3\n",
    "def download_metrics():\n",
    "    # Get the job ID from the ARN\n",
    "    job_id = job_arn.split('/')[-1]\n",
    "    \n",
    "    # Define file paths\n",
    "    train_metrics_s3_key = f\"output/model-customization-job-{job_id}/training_artifacts/step_wise_training_metrics.csv\"\n",
    "    \n",
    "    local_train_metrics = \"train_metrics.csv\"\n",
    "    \n",
    "    # Download files\n",
    "    try:\n",
    "        s3_client.download_file(bucket_name, train_metrics_s3_key, local_train_metrics)\n",
    "        \n",
    "        print(\"Metrics downloaded successfully\")\n",
    "        return local_train_metrics\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"Error downloading metrics: {e}\")\n",
    "        return None, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b031e9a4-762b-41c5-a4d9-c3e278050299",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download metrics\n",
    "train_metrics_file = download_metrics()\n",
    "\n",
    "# Plot training and validation loss if metrics are available\n",
    "if train_metrics_file:\n",
    "    import pandas as pd\n",
    "    \n",
    "    # Load metrics\n",
    "    train_data = pd.read_csv(train_metrics_file)\n",
    "\n",
    "    # Calculate step-level training loss\n",
    "    train_metrics_epoch = train_data.groupby('step_number').mean()\n",
    "    \n",
    "    # Plot\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    plt.plot(train_metrics_epoch.index, train_metrics_epoch.training_loss, label='Training')\n",
    "    plt.title('Training Loss')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.xlabel('Step')\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb4aebb3-05de-46c6-92cb-661dcf348384",
   "metadata": {},
   "source": [
    "## Create Provisioned Throughput\n",
    "\n",
    "To use the fine-tuned model for inference, we need to create provisioned throughput:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b67741f-632c-4927-911a-37095f171322",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a unique name for provisioned throughput\n",
    "provisioned_model_name = f\"llama32-multi-model-prov-{timestamp}\"\n",
    "\n",
    "# Create provisioned throughput\n",
    "try:\n",
    "    response = bedrock.create_provisioned_model_throughput(\n",
    "        modelId=custom_model_arn,\n",
    "        provisionedModelName=provisioned_model_name,\n",
    "        modelUnits=1  \n",
    "    )\n",
    "    \n",
    "    provisioned_model_arn = response[\"provisionedModelArn\"]\n",
    "    print(f\"Provisioned model created: {provisioned_model_arn}\")\n",
    "    \n",
    "    # Monitor provisioning status\n",
    "    status = bedrock.get_provisioned_model_throughput(provisionedModelId=provisioned_model_arn)[\"status\"]\n",
    "    print(f\"Initial provisioning status: {status}\")\n",
    "    \n",
    "    progress_bar = tqdm(desc=\"Provisioning model\", bar_format=\"{desc}: {bar}\")\n",
    "    \n",
    "    # Poll provisioning status\n",
    "    while status == \"Creating\":\n",
    "        time.sleep(60)  # Check every minute\n",
    "        status = bedrock.get_provisioned_model_throughput(provisionedModelId=provisioned_model_arn)[\"status\"]\n",
    "        progress_bar.update(1)\n",
    "        progress_bar.set_description(f\"Provisioning model (Status: {status})\")\n",
    "    \n",
    "    progress_bar.close()\n",
    "    \n",
    "    # Final status\n",
    "    final_status = bedrock.get_provisioned_model_throughput(provisionedModelId=provisioned_model_arn)[\"status\"]\n",
    "    print(f\"Final provisioning status: {final_status}\")\n",
    "    \n",
    "except Exception as e:\n",
    "    print(f\"Error creating provisioned throughput: {e}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dc29c55-1ef4-460d-9fc5-923bdf958406",
   "metadata": {},
   "source": [
    "## Test with Inference\n",
    "\n",
    "Now, let's test our fine-tuned model using the test dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04100116-5267-43f4-ace6-68e35891fd02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to convert S3 URI to local image path\n",
    "def s3_uri_to_local_path(s3_uri):\n",
    "    \"\"\"Convert S3 URI to local file path\"\"\"\n",
    "    filename = s3_uri.split('/')[-1]\n",
    "    fileformat = filename.split('.')[-1]\n",
    "    docname = filename.split('.')[0]\n",
    "\n",
    "    if fileformat != \"png\":\n",
    "        img = Image.open(f\"llava_images/test/{filename}\")\n",
    "        img.save(f\"llava_images/test/{docname}.png\", optimize=True, compress_level=9)\n",
    "        return f\"llava_images/test/{docname}.png\"\n",
    "\n",
    "    return f\"llava_images/test/{filename}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47503db0-d10f-4d54-a1c5-b8a5cb6ea3ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to run inference using Bedrock converse API\n",
    "def run_inference(image_path, question, model_id):\n",
    "    \"\"\"Run inference using the converse API with local image\"\"\"\n",
    "    \n",
    "    # Read image as binary data directly\n",
    "    with open(image_path, \"rb\") as f:\n",
    "        image_bytes = f.read()\n",
    "\n",
    "    # Create message structure matching the example\n",
    "    message = {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\n",
    "                \"image\": {\n",
    "                    \"format\": \"png\",  # Our images are jpg format\n",
    "                    \"source\": {\n",
    "                        \"bytes\": image_bytes  # Raw bytes, no base64 encoding\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            {\n",
    "                \"text\": question\n",
    "            }\n",
    "        ]\n",
    "    }\n",
    "\n",
    "    inference_config = {\"temperature\": 0.01}\n",
    "    \n",
    "    # Call the converse API\n",
    "    response = bedrock_runtime.converse(\n",
    "        modelId=model_id,\n",
    "        messages=[message],\n",
    "        inferenceConfig=inference_config\n",
    "    )\n",
    "    \n",
    "    # Extract response text following the example\n",
    "    response_text = response[\"output\"][\"message\"][\"content\"][0][\"text\"]\n",
    "    return response_text\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62f126a8-8304-433d-846d-be9868781a54",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load test samples from local file\n",
    "with open('test.jsonl', 'r') as f:\n",
    "    test_samples = [json.loads(line) for line in f.readlines()][:5]  # Get first 5 samples\n",
    "\n",
    "# Run inference on test samples\n",
    "for i, sample in enumerate(test_samples):\n",
    "    # Extract information\n",
    "    question = sample[\"messages\"][0][\"content\"][0][\"text\"]\n",
    "    image_s3_uri = sample[\"messages\"][0][\"content\"][1][\"image\"][\"source\"][\"s3Location\"][\"uri\"]\n",
    "    expected_answer = sample[\"messages\"][1][\"content\"][0][\"text\"]\n",
    "    \n",
    "    # Convert S3 URI to local file path\n",
    "    local_image_path = s3_uri_to_local_path(image_s3_uri)\n",
    "    \n",
    "    print(f\"\\n=== Test Sample {i+1} ===\")\n",
    "    print(f\"Question: {question}\")\n",
    "    print(f\"Expected Answer: {expected_answer}\")\n",
    "    \n",
    "    try:\n",
    "        # Run inference\n",
    "        model_response = run_inference(local_image_path, question, provisioned_model_arn)\n",
    "        print(f\"Model Response: {model_response}\")\n",
    "    except Exception as e:\n",
    "        print(f\"Inference error: {e}\")\n",
    "        # If the provisioned model ARN doesn't work, try using the regular model ID format\n",
    "        try:\n",
    "            # Try with a standard model ID format as shown in the example\n",
    "            model_response = run_inference(local_image_path, question, provisioned_model_arn)\n",
    "            print(f\"Model Response (using standard model ID): {model_response}\")\n",
    "        except Exception as e2:\n",
    "            print(f\"Second attempt failed: {e2}\")\n",
    "    \n",
    "    # Display the image\n",
    "    image = Image.open(local_image_path)\n",
    "    plt.figure(figsize=(6, 6))\n",
    "    plt.imshow(image)\n",
    "    plt.axis('off')\n",
    "    plt.title(f\"Test Image {i+1}\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "804b1a9e-8db3-4665-b146-8879861449fe",
   "metadata": {},
   "source": [
    "## Clean Up Resources\n",
    "\n",
    "Finally, let's clean up the resources we created:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1983d9f6-7628-4f82-b617-402c80d0bc91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to clean up resources\n",
    "def clean_up():\n",
    "    print(\"Cleaning up resources...\")\n",
    "    \n",
    "    # Delete provisioned model throughput\n",
    "    try:\n",
    "        print(\"Deleting provisioned model throughput...\")\n",
    "        bedrock.delete_provisioned_model_throughput(provisionedModelId=provisioned_model_arn)\n",
    "        print(\"Provisioned model throughput deleted\")\n",
    "    except Exception as e:\n",
    "        print(f\"Error deleting provisioned model throughput: {e}\")\n",
    "    \n",
    "    # Clean up IAM resources\n",
    "    iam = session.client('iam')\n",
    "    try:\n",
    "        print(\"Detaching policy from role...\")\n",
    "        iam.detach_role_policy(RoleName=role_name, PolicyArn=policy_arn)\n",
    "        \n",
    "        print(\"Deleting policy...\")\n",
    "        iam.delete_policy(PolicyArn=policy_arn)\n",
    "        \n",
    "        print(\"Deleting role...\")\n",
    "        iam.delete_role(RoleName=role_name)\n",
    "        \n",
    "        print(\"IAM resources cleaned up\")\n",
    "    except Exception as e:\n",
    "        print(f\"Error cleaning up IAM resources: {e}\")\n",
    "    \n",
    "    # We're not deleting the S3 bucket here as you might want to keep your data and model\n",
    "    \n",
    "    print(\"Cleanup completed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c59940a-11b9-494a-9c38-533df730478d",
   "metadata": {},
   "outputs": [],
   "source": [
    "clean_up()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ec4d8e9-f1ee-4173-bc6b-00e38ddf5d75",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this notebook, we successfully fine-tuned a Llama 3.2 multi-modal model using Amazon Bedrock. We:\n",
    "\n",
    "- Set up and launched a fine-tuning job with our prepared dataset\n",
    "- Monitored the job progress and visualized training metrics\n",
    "- Created provisioned throughput for the fine-tuned model\n",
    "- Tested the model's performance with inference on test samples\n",
    "- Cleaned up resources we no longer needed\n",
    "\n",
    "The fine-tuned model can now answer questions about images based on the patterns it learned from our training data. For real-world applications, you may want to use a larger and more diverse dataset tailored to your specific use case."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
